import os,whisper,re,datetime,zhconv,jieba,pickle
from moviepy import VideoFileClip



def get_allvideopath(path,txtfile):
    allfilepath = ''
    for root, dirs, files in os.walk(path):
        for file in files:
            if not file.startswith('.') and file.lower().endswith(('.mp4','.MP4', '.MOV','.ts')):
                allfilepath+=os.path.join(root, file)+'\n'
    open(txtfile,'w',encoding='utf-8').write(allfilepath)

def video_2_mp3(txt_allvideopath,output_folder):
    # 检查并创建输出文件夹
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 遍历文件夹中的所有文件
    for filename in open(txt_allvideopath,'r',encoding='utf-8').read().split('\n'):

        output_path = output_folder + '/'+ filename.split('\\')[-1].replace('.mp4','.mp3')
        if os.path.exists(output_path):
            print('ok',end=',')
            continue
        
        try:
            # 打开视频文件
            video = VideoFileClip(filename)
            
            # 检查是否有音频轨道
            if video.audio:
                # 提取音频并保存
                audio = video.audio
                audio.write_audiofile(output_path)
                audio.close()

            # else:
            #     print(f"No audio track in {filename}")
            
            # 关闭视频文件
            video.close()
        except Exception as e:
            print(f"Error processing {filename}: {e}")

def transcribe_audio(file_path,timeonoff):
    # whisper mp3 return str
    device = "cuda"  # 使用 GPU
    model = whisper.load_model("turbo",device=device)  # 加载模型
    result = model.transcribe(file_path, language="zh")  # 识别音频
    # 输出带时间戳的转录结果
    text_i = ''
    for segment in result["segments"]:
    
        start_time = segment["start"]
        start_time = datetime.datetime.fromtimestamp(start_time,tz=datetime.timezone.utc).strftime("%H:%M:%S")

        if timeonoff=='on':
            text_i += f"[{start_time}] {segment["text"]}\n"
        else:
            text_i += f"{segment["text"]}\n"
    return text_i
def mp3_whisper_to_text(output_txt_folder,output_mp3_folder,timeonoff='on'):
    # 检查并创建输出文件夹
    if not os.path.exists(output_txt_folder):
        os.makedirs(output_txt_folder)

    print('%s/%s is finished' %(len(os.listdir(output_txt_folder)),len(os.listdir(output_mp3_folder))))
    if not os.path.exists(output_txt_folder):
        os.makedirs(output_txt_folder)
    for cd_i in sorted(os.listdir(output_mp3_folder)):
        if os.path.exists(os.path.join(output_txt_folder,cd_i.replace('.mp3','.txt'))):
            continue
        text_i = transcribe_audio(os.path.join(output_mp3_folder,cd_i),timeonoff)
        text_i = zhconv.convert(text_i, "zh-cn")
        open(os.path.join(output_txt_folder,cd_i.replace('.mp3','.txt')),'w',encoding='utf-8').write(text_i)
        print(cd_i,'..ok')
        
# 构建倒排索引
def build_inverted_index(directory):
    inverted_index = {}
    for filename in os.listdir(directory):
        filename_fix = filename.replace('-',"/").replace('.txt',"")
        if filename.endswith(".txt"):
            file_path = os.path.join(directory, filename)
            with open(file_path, 'r', encoding='utf-8') as file:
                lines = file.readlines()
                for line in lines:
                    line = line.strip()
                    # 使用正则表达式提取时间和内容
                    match = re.match(r'^\[(.*?)\]\s*(.*)$', line)
                    if match:
                        time_str = match.group(1)
                        content = match.group(2)
                        # 使用 jieba 进行分词
                        words = jieba.lcut(content)
                        # 过滤停用词
                        words = [word for word in words if word not in stop_words]
                        for word in words:
                            key = word
                            value = f"{content}\n        -> [{filename_fix}][{time_str}]"
                            if key in inverted_index:
                                if value not in inverted_index[key]:
                                    inverted_index[key].append(value)
                            else:
                                inverted_index[key] = [value]
    return inverted_index
# 查询倒排索引
def query_index(inverted_index, query):
    # 尝试将查询词作为一个整体进行匹配
    if query in inverted_index:
        return inverted_index[query]
    else:
        isbreakwords = input("整体匹配失败，输入[1]尝试单独的词进行匹配,其他则放弃")
        if isbreakwords == '1':
            query_words = jieba.lcut(query)
            results = set()
            for word in query_words:
                if word in inverted_index:
                    results.update(inverted_index[word])
            return list(results)

def check_inverted_words(directory,index_file):
    if input('输入[1],清理之前的倒排索引重新生成,任意键继续')=='1':
        print(f'delete {index_file}')
        os.remove(index_file)
    if os.path.exists(index_file):
        with open(index_file, 'rb') as f:
            inverted_index = pickle.load(f)
        print("倒排索引加载完成！")
    else:
        print("正在构建倒排索引，请稍候...")
        inverted_index = build_inverted_index(directory)
        with open(index_file, 'wb') as f:
            pickle.dump(inverted_index, f)
        print("倒排索引构建并保存完成！")
    
    while True:
        query = input("请输入要查询的词语（输入'end'结束）：")
        if query == 'end':
            break
        results = query_index(inverted_index, query)
        if results:
            print(f"\n包含词语 '{query}' 的文件和时间戳有：")
            for entry in results:
                print(entry)
        else:
            print(f"没有找到包含词语 '{query}' 的文件。")


path_all = 'D:\\00_python_auto\\video2mp32txt\\00-orivideo'
work_dir = 'D:\\00_python_auto\\video2mp32txt'
txt_allvideopath = os.path.join(work_dir,'all_file_path.txt')
output_mp3_folder = os.path.join(work_dir,'audios')
output_txt_folder = os.path.join(work_dir,'txt')
index_file = os.path.join(work_dir,'inverted_index.pkl')
# 定义停用词列表
stop_words = set(open(os.path.join(work_dir,'cn_stopwords.txt'),'r',encoding='utf_8_sig').read().split('\n'))
# 自定义词条
# for words_i in ["500强",]:
#     jieba.add_word(words_i)


# print('2-将所有视频文件转化为MP3')
# get_allvideopath(path_all,txt_allvideopath)
# video_2_mp3(txt_allvideopath,output_mp3_folder)

print('3-将所有音频文件转化为txt')
mp3_whisper_to_text(output_txt_folder,output_mp3_folder,timeonoff='off')

# print('4-构建倒排索引，查询倒排索引，返回词语的文件和时间戳有')
# check_inverted_words(output_txt_folder,index_file)

# print('5-将所有txt汇总成为一个txt')
# onetxt= ''
# alltxtinone = os.path.join(work_dir,'alltxtinone.txt')

# open(alltxtinone,'w',encoding='utf_8_sig').write('')
# for cd_i in os.listdir('video2mp32txt/txt'):
#     # onetxt = cd_i +'\n'
#     onetxt= ''
#     onetxt += open('video2mp32txt/txt/'+cd_i,'r',encoding='utf-8').read()+'\n'
#     open(alltxtinone,'a+',encoding='utf_8_sig').write(onetxt)








